# data handling and visualisation
library(tidyverse)
library(lubridate)
library(ggpubr)

# FPCA + LMEs
library(fdapace)
library(lme4)
library(lmerTest)
library(emmeans)

# GAMMs
library(mgcv)
library(itsadug)
library(tidygam) # for 'get_gam_predictions'
source("https://raw.githubusercontent.com/soskuthy/gamm_intro/master/gamm_hacks.r")

## loading data
# define the path
dir <- getwd()

# index wav files in the directory
file_list <- list.files(paste0(dir, "/ultrasound/spline"), pattern = ".txt", full.names = FALSE)

# create an empty list to store data
data_list <- list()

for(i in 1:length(file_list)){
  current_data <- read.delim(paste0(dir, "/ultrasound/spline/", file_list[i]), header = FALSE)
  
  data_list[[i]] <- current_data
}

# bind all data from the list into a data frame
dat <- bind_rows(data_list)

# omit na
dat <- na.omit(dat)

dat <- dat |> 
  janitor::remove_empty(which = "cols") |>  # remove empty columns
  dplyr::rename(
    speaker = V1, # speaker ID
    rec_date = V2, # date/time associated with each recording
    prompt = V3, # alias of prompt and repetition
    interval = V4 # textgrid label
    ) |> 
  dplyr::rename(
    X_1 = V5, # vallecula
    Y_1 = V6,
    X_2 = V7, # tongue root 1
    Y_2 = V8,
    X_3 = V9, # tongue root 2
    Y_3 = V10,
    X_4 = V11, # tongue body 1
    Y_4 = V12,
    X_5 = V13, # tongue body 2
    Y_5 = V14,
    X_6 = V15, # tongue dorsum 1
    Y_6 = V16,
    X_7 = V17, # tongue dorsum 2
    Y_7 = V18,
    X_8 = V19, # tongue blade 1
    Y_8 = V20,
    X_9 = V21, # tongue blade 2
    Y_9 = V22,
    X_10 = V23, # tongue tip 1
    Y_10 = V24,
    X_11 = V25, # tongue tip 2
    Y_11 = V26
  )

# eliminate one token that was not labelled properly
dat <- dat |> 
  filter(!str_detect(rec_date, "2022/11/29 17:48:39")
  )

# Specifying which time point of 11 each tongue spline is associated with. 
dat <- dat |> 
  group_by(rec_date, speaker) |> 
  mutate(prop_time = row_number() - 1) |>  # so that first row = 0
  ungroup() |> 
  mutate(prop_time = prop_time * 10) |>  # to give percentages
  mutate(prop_time = factor(prop_time)) |> 
  mutate(perc_time = paste(prop_time, "%", sep = "") |>  factor())

# Rename the prompt '18' into 'biribiri' as AAA wasn't able to handle the Japanese characters.
dat$prompt[dat$prompt == "18"] <- "biribiri"

# check data
dat

## pivot longer: Use this when you work on the tongue only
dat_long <- dat |> 
  pivot_longer(5:26, names_to = c(".value", "point_number"), names_sep = "_")

# check 
dat_long

## 1. Getting R to recognise dates and time with the lubridate package
dat_long$rec_date <- lubridate::parse_date_time(dat_long$rec_date, c("Ymd HMS","dmY HM"))
# multiple formats are specified here because some rec_dates look weird when exported from AAA

## 2. create a summary of the data using summarise(), add repetition information (as it's easier) and store it in the object 'rep' 
rep <- dat_long |> 
  dplyr::group_by(speaker, prompt, rec_date) |> 
  dplyr::summarise() |> 
  dplyr::mutate(
    repetition = row_number()
  ) |> 
  ungroup()


## 3. merge 'rep' with the main data ('df.long' here) using 'left_join()'
dat_long <- dplyr::left_join(dat_long, rep, by = c("speaker", "prompt", "rec_date"))

dat_long <- dat_long |> 
  filter(!speaker == "2d57ke" | !repetition == "1")

dat_long <- dat_long |> 
  dplyr::mutate(
    token_id = paste(speaker, prompt, sep = "_"),
    token_id = paste(token_id, repetition, sep = "_")
  )

### a combination of "rec_date" and "prop_time" gives unique label for each spline
dat_long <- dat_long |> 
  dplyr::mutate(
    spline_id = str_c(dat_long$token_id, dat_long$prop_time, sep = "_")
  )

## Z score normalisation of X Y coordinates
dat_long <- dat_long |> 
  dplyr::group_by(speaker) |> 
  dplyr::mutate(
    X_z = as.numeric(scale(X)),
    Y_z = as.numeric(scale(Y))
  ) |> 
  dplyr::ungroup()

# perform PCA
pca_results <- prcomp(dplyr::select(dat, 5:26), scale = FALSE)

# summary
summary(pca_results)

# Check the object structure
str(pca_results)

# Mean values from the output of the PCA
pca_mean <- tibble::enframe(pca_results$center)

# subset data to make into a matrix of x and y values
pca_mean <- pca_mean |> 
  dplyr::mutate(
    axis_mean = str_sub(name, start = 1, end = 1),
    number_mean = str_sub(name, start = 3, end = -1)
  ) |> 
  dplyr::select(-name) |> 
  dplyr::relocate(axis_mean, number_mean)

# extract X coordinate values of the mean values
mean_X <- pca_mean |> 
  dplyr::filter(axis_mean == "X") |> 
  dplyr::rename(mean_X = value,
                axis_mean_X = axis_mean,
                number_mean_X = number_mean)

# extract Y coordinate values of the mean values
mean_Y <- pca_mean |> 
  dplyr::filter(axis_mean == "Y") |> 
  dplyr::rename(mean_Y = value,
                axis_mean_Y = axis_mean,
                number_mean_Y = number_mean)

# combine X and Y values together
mean_pca <- cbind(mean_X, mean_Y)

## get loadings - eigenvectors
loadings <- as.table(pca_results$rotation)

# get sds of first 2 PCs
sd <- tibble::enframe(pca_results$sdev)
sd_PC1 <- as.numeric(sd[1,2])
sd_PC2 <- as.numeric(sd[2,2])

# calculate estimated values including sd
## PC1
estimate_PC1 <- cbind(mean_pca, PC1_loadings)
estimate_PC1$PC1_max_X <- estimate_PC1$mean_X + sd_PC1*estimate_PC1$PC1_loadings_X
estimate_PC1$PC1_min_X <- estimate_PC1$mean_X - sd_PC1*estimate_PC1$PC1_loadings_X
estimate_PC1$PC1_max_Y <- estimate_PC1$mean_Y + sd_PC1*estimate_PC1$PC1_loadings_Y
estimate_PC1$PC1_min_Y <- estimate_PC1$mean_Y - sd_PC1*estimate_PC1$PC1_loadings_Y

# calculate estimated values including sd
## PC2
estimate_PC2 <- cbind(mean_pca, PC2_loadings)
estimate_PC2$PC2_max_X <- estimate_PC2$mean_X + sd_PC2*estimate_PC2$PC2_loadings_X
estimate_PC2$PC2_min_X <- estimate_PC2$mean_X - sd_PC2*estimate_PC2$PC2_loadings_X
estimate_PC2$PC2_max_Y <- estimate_PC2$mean_Y + sd_PC2*estimate_PC2$PC2_loadings_Y
estimate_PC2$PC2_min_Y <- estimate_PC2$mean_Y - sd_PC2*estimate_PC2$PC2_loadings_Y

# PC1
PC1_reconstructed <- ggplot() +
  geom_line(data = estimate_PC1, aes(x = mean_X, y = mean_Y), size = 1.5) +
  geom_line(data = estimate_PC1, aes(x = PC1_max_X, y = PC1_max_Y), linetype = "dashed", linewidth = 1, alpha = 0.5, colour = "#E69F00") +
  geom_line(data = estimate_PC1, aes(x = PC1_min_X, y = PC1_min_Y), linetype = "dotted", linewidth = 1, alpha = 0.5, colour = "#56B4E9") +
  geom_point(data = estimate_PC1, aes(x = PC1_max_X, y = PC1_max_Y), shape = 3, size = 3, stroke = 2, colour = "#E69F00") +
  geom_point(data = estimate_PC1, aes(x = PC1_min_X, y = PC1_min_Y), shape = "\u2212", size = 5, stroke = 8, colour = "#56B4E9") +
  labs(x = "X", y = "Y", title = "PC1")

# PC2
PC2_reconstructed <- ggplot() +
  geom_line(data = estimate_PC2, aes(x = mean_X, y = mean_Y), size = 1.5) +
  geom_line(data = estimate_PC2, aes(x = PC2_max_X, y = PC2_max_Y), linetype = "dashed", linewidth = 1, alpha = 0.5, colour = "#E69F00") +
  geom_line(data = estimate_PC2, aes(x = PC2_min_X, y = PC2_min_Y), linetype = "dotted", linewidth = 1, alpha = 0.5, colour = "#56B4E9") +
  geom_point(data = estimate_PC2, aes(x = PC2_max_X, y = PC2_max_Y), shape = 3, size = 3, stroke = 2, colour = "#E69F00") +
  geom_point(data = estimate_PC2, aes(x = PC2_min_X, y = PC2_min_Y), shape = "\u2212", size = 5, stroke = 8, colour = "#56B4E9") +
  labs(x = "X", y = "Y", title = "PC2")

# showing plots
ggpubr::ggarrange(PC1_reconstructed, PC2_reconstructed, nrow = 1)

# Check the attributes again
str(pca_results)

# Get the results of the PCA which are useful
pc_score <- pca_results$x

# Put it into a sensible format as some variables come out weird
pc_score <- as_tibble(pc_score) |> 
  dplyr::select(1:4)

# Combine with non-numeric information from earlier
dat_pc <- cbind(dat, pc_score)

# normalise PCs by speaker for comparison
pca_result <- dat_pc |> 
  dplyr::group_by(speaker) |>  
  dplyr::mutate(
    PC1z = scale(PC1),
    PC2z = scale(PC2),
    PC3z = scale(PC3),
    PC4z = scale(PC4)
  ) |> 
  dplyr::ungroup()

# converting data frame into long format (and parsing rec_date)
pca_result_long <- pca_result |> 
  tidyr::pivot_longer(5:26, names_to = c(".value", "point_number"), names_sep = "_") |> 
  dplyr::mutate(
    rec_date = lubridate::parse_date_time(rec_date, c("Ymd HMS","dmY HM"))
  )

# integrate meta data by joining pca_result_long to dat_long
dat_long <- dplyr::left_join(dat_long, pca_result_long, by = c("speaker", "rec_date", "prompt", "interval", "prop_time", "perc_time", "point_number", "X", "Y"))

# add L1 information
dat_long <- dat_long |> 
  dplyr::mutate(
    L1 = case_when(
      speaker %in% c("4ps8zx", "5jzj2h", "5upwe3", "6p63jy", "cu2jce", "ds6umh", "h5s4x3", "jcy8xi", "m46dhf", "tay55n", "we8z58", "xub9bc") ~ "English",
      TRUE ~ "Japanese"
    )
  )

# IDs = token column; tVec = time column; yVec = variable column(s)
input_dat <- fdapace::MakeFPCAInputs(IDs = dat_sub$token_id, tVec = dat_sub$prop_time, yVec = dat_sub$PC1z)

# Check if there's any issues with the data
fdapace::CheckData(input_dat$Ly, input_dat$Lt)

# No errors have been returned, so let's now run FPCA 
dat_dyn_fpca <- fdapace::FPCA(Ly = input_dat$Ly, Lt = input_dat$Lt, optns = list(plot = TRUE))

dat_dyn_fpca$cumFVE

# function: get PC scores + return data frame with PCs for each token
get_pc_scores <- function(fpcaObj){
  pcs <- data.frame(fpcaObj$xiEst)
  token <- names(fpcaObj$inputData$Lt) 
  df <- cbind(token, pcs)
  n_pcs <- length(fpcaObj$lambda) # get number of PCs
  pc_names <- paste0("PC", 1:n_pcs) # create colnames for PCs
  names(df) <- c("token_id", pc_names) # add colnames for token + PCs
  return(df)
}

# get PC scores w/ token info
pc_df <- get_pc_scores(dat_dyn_fpca)

# join PCs (dat) with selected cols from original data frame 
## store meta info
meta <- dat_sub |>  
  dplyr::select(speaker, L1, rec_date, prompt, interval, prop_time, repetition, token_id)

## merge the list and meta data - unique(meta) because otherwise there would be lots of duplicates
dat <- left_join(pc_df, unique(meta), by = "token_id")

# function: define perturbation function (±Q = ±sd, k = PC number)
perturbation <- function(fpcaObj, Q, k){
  Q * sqrt(fpcaObj$lambda[k]) * fpcaObj$phi[,k] + fpcaObj$mu
}

# function: create perturbation object with mean and ±Q sd as a data frame (for one PC only)
# can validate against fdapace::GetMeanCurve and fdapace::CreateModeOfVarPlot
perturbation_object <- function(fpcaObj, Q, k){
  time <- fpcaObj$workGrid # grid of time values
  mean <- fpcaObj$mu # mean trajectory
  Qplus <- perturbation(fpcaObj, Q, k) # +Q sd
  Qminus <- perturbation(fpcaObj, -Q, k) # -Q sd
  df <- cbind(time, mean, Qplus, Qminus)
  colnames(df) <- c("time", "mean", "Qplus", "Qminus")
  df <- data.frame(df)
  df$PC <- paste0("PC", k) # add PC colname
  return(df)
}

# function: create perturbation data frame with mean and ±Q sd (for all PCs)
# to do: add ability to pass list of Q values for gradient perturbation function
get_perturbation <- function(fpcaObj, Q){
  n_pcs <- length(fpcaObj$lambda)
  k <- 1:n_pcs
  df <- lapply(k, perturbation_object, fpcaObj=fpcaObj, Q=Q)
  df <- dplyr::bind_rows(df) # unnest lists into single df
  return(df)
}

# get mean trajectory and ±2 sd for all PCs
Q <- seq(from = -2, to = 2, by = 0.1)
pQ <- lapply(Q, get_perturbation, fpcaObj = dat_dyn_fpca) # change fpcaObj appropriately
names(pQ) <- Q # name pQ lists using value of Q
pQ <- dplyr::bind_rows(pQ, .id = "Q") # collapse lists together
pQ$Q <- as.numeric(pQ$Q) # make 'Q' column numeric

# visualise variation along each FPC with mean function
pQ <- pQ |> 
  dplyr::filter(PC %in% c('PC1','PC2', 'PC3')) |> 
  dplyr::mutate(
    PC = case_when(
      PC == "PC1" ~ "FPC1",
      PC == "PC2" ~ "FPC2",
      PC == "PC3" ~ "FPC3"
    )
  ) |> 
  tidyr::pivot_longer(mean:Qminus, names_to = "Qsd", values_to = "value")

pQ |>
  ggplot2::ggplot() +
  aes(x = time, y = value, colour = Q, group = interaction(Q, Qsd)) +
  geom_path() +
  geom_line(data = pQ |>  dplyr::filter(Qsd == "mean"), aes(x = time, y = value), colour = "black", size = 1) + 
  facet_wrap(~ PC, ncol = 2) +
  scale_colour_gradient2(low = "#E69F00", mid = "#56B4E9", high = "#009E73", midpoint = 0)+
  labs(x = "Time (normalised)", y = "Reconstructed PC1 trajectories", color = "FPC score")

# mean fPC1 trajectory
# pc1_mean_curve <- fdapace::GetMeanCurve(Ly = input.PC1$Ly, Lt = input.PC1$Lt, optns = list(plot = TRUE))
mu_values <- data.frame(dat_dyn_fpca$mu) # mean curve values
mu_time <- data.frame(dat_dyn_fpca$workGrid) # timepoints used for estimating the curve
phi <- data.frame(dat_dyn_fpca$phi) # eigenfunction at each timepoint: workGrid * nlambda (e.g., 255 = 51 workGrid * 5 lambda)
lambda <- data.frame(dat_dyn_fpca$lambda) # PC loadings for each PC: currently 5

# create a data frame containing mean curve, time and eigenfunctions assocaited with each PC at each time point
## add an extra column 'col_number' as a common index across the data frames - useful when merging everything together later on
### mean curve
mu_values <- mu_values |>  
  dplyr::mutate(
    col_number = row_number()
  )

### sampling time points
mu_time <- mu_time |>  
  dplyr::mutate(
    col_number = row_number()
  )

### eigenfunction
phi <- phi |>  
  dplyr::mutate(
    col_number = row_number()
  )

### pc loadings
lambda <- lambda |>  
  dplyr::mutate(
    PC = str_c("PC", row_number()),
    PC = str_c(PC, "lambda", sep = "_")
  ) |>  
  tidyr::pivot_wider(names_from = "PC", values_from = "dat_dyn_fpca.lambda") |>  
  dplyr::slice(rep(1:n(), each = 11)) |>  
  dplyr::mutate(
    col_number = row_number()
  )
  
## merging all data together one by one
rec <- dplyr::left_join(mu_values, mu_time, by = "col_number")
rec <- dplyr::left_join(rec, phi, by = "col_number")
rec <- dplyr::left_join(rec, lambda, by = "col_number")

## tidying up some column names
rec <- rec |>  
  dplyr::select(col_number, dat_dyn_fpca.workGrid, dat_dyn_fpca.mu, X1, X2, X3, X4, PC1_lambda, PC2_lambda, PC3_lambda, PC4_lambda) |>  
  dplyr::rename(
    mean = dat_dyn_fpca.mu,
    time = dat_dyn_fpca.workGrid,
    PC1_eigen = X1,
    PC2_eigen = X2,
    PC3_eigen = X3,
    PC4_eigen = X4
  )

## plotting the eigenfunctions - this should match with a sub-plot in bottom right created with plot(PC1)
rec |>  
  ggplot() +
  # geom_path(aes(x = time, y = mean)) +
  geom_path(aes(x = time, y = PC1_eigen), colour = "black", linewidth = 1.5) +
  geom_path(aes(x = time, y = PC2_eigen), colour = "red", linetype = 2, linewidth = 1.5) +
  geom_path(aes(x = time, y = PC3_eigen), colour = "darkgreen", linetype = 3, linewidth = 1.5) +
  # geom_path(aes(x = time, y = value, colour = pc)) +
  geom_hline(yintercept = 0) +
  labs(x = "time", y = "eigenfunctions", title = "First 3 eigenfunctions")

## check if this matches plot(dat_dyn_fpca)
plot(dat_dyn_fpca)

# PC scores -> each row is 1 token, each column is one PC
head(dat_dyn_fpca$xiEst)

# PC scores have already been added to the main data set
head(dat)

# duplicate each row by 11 times 
dat_time <- dat |> 
  dplyr::slice(rep(1:n(), each = 11))

# add col_names to merge with the other data frame
dat_time <- dat_time |>  
  dplyr::group_by(token_id) |>  
  dplyr::mutate(
    col_number = row_number()
  ) |>  
  ungroup()

# merge
dat_time <- left_join(dat_time, rec, by = "col_number") |> 
  na.omit()

# reconstruct individual trajectory tokens
dat_time <- dat_time |>  
  dplyr::mutate(
    PC1_reconstruct = PC1 * PC1_eigen + mean,
    PC2_reconstruct = PC2 * PC2_eigen + mean,
    PC3_reconstruct = PC3 * PC3_eigen + mean,
    PC4_reconstruct = PC4 * PC4_eigen + mean
  )

# visualise FPC1
dat_time |>  
  ggplot() +
  geom_path(aes(x = time, y = PC1_reconstruct, group = token_id, colour = prompt), alpha = 0.2, show.legend = TRUE) +
  scale_color_manual(values = c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")) +
  labs(x = "Proportional Time (%)", y = "Reconstructed PC1", title = "Reconstructed PC1 teajectories from FPC1") +
  geom_hline(yintercept = 0, linetype = 1, linewidth = 0.1) +
  guides(colour = guide_legend(override.aes = list(alpha = 1))) +
  facet_wrap(~ L1)

# visualise FPC2
dat_time |>  
  ggplot() +
  geom_path(aes(x = time, y = PC2_reconstruct, group = token_id, colour = prompt), alpha = 0.2, show.legend = TRUE) +
  scale_color_manual(values = c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")) +
  labs(x = "Proportional Time (%)", y = "Reconstructed PC1", title = "Reconstructed PC1 teajectories from FPC2") +
  geom_hline(yintercept = 0, linetype = 1, linewidth = 0.1) +
  guides(colour = guide_legend(override.aes = list(alpha = 1))) +
  facet_wrap(~ L1)

# check data
colnames(dat_sub)

# add 'segment', convert variables into factor and create an interaction variable
dat_sub <- dat_sub |> 
  dplyr::mutate(
    segment = case_when(
      interval == "ili" ~ "/l/",
      interval == "iri" ~ "/ɹ/",
      interval == "idi" ~ "/ɾ/"
    ),
    segment = as.factor(segment),
    L1 = as.factor(L1),
    speaker = as.factor(speaker),
    prompt = as.factor(prompt),
    LangSeg = interaction(L1, segment)
  )

# also make sure that prop_time is numeric
dat_sub <- dat_sub |> 
  dplyr::mutate(
    prop_time = as.numeric(prop_time)
  )

# converting LangCont into ordered variable
dat_sub$LangSeg.ord <- as.ordered(dat_sub$LangSeg)
contrasts(dat_sub$LangSeg.ord) <- "contr.treatment" # don't forget this!

# fitting a model
m1 <- mgcv::bam(PC1z ~ LangSeg.ord + s(prop_time) + s(prop_time, by = LangSeg.ord) + s(prop_time, speaker, bs = "fs", m = 1), data = dat_sub, method = "fREML")

# summary
summary(m1)

# mark start.event
## arrange the data appropriately
dat_sub <- dat_sub |> 
  dplyr::arrange(token_id, prop_time)

## mark time = 1 (as we express prop_time between 1 and 11)
dat_sub$start.event <- dat_sub$prop_time == 1

# check residual autocorrelations
itsadug::acf_resid(m1)

# obtain the autocorrelation at lag 1 as a rho value
rho_m1 <- start_value_rho(m1) # 0.7551264

# fitting a model
m2 <- mgcv::bam(PC1z ~ LangSeg.ord + s(prop_time) + s(prop_time, by = LangSeg.ord) + s(prop_time, speaker, bs = "fs", m = 1), data = dat_sub, method = "fREML", rho = rho_m1, AR.start = dat_sub$start.event)

# summary
summary(m2)

# check residual autocorrelations -- looks fairly good
itsadug::acf_resid(m2)